import cv2
import numpy as np


def extend_border(image, top, bottom, left, right, border_type=cv2.BORDER_CONSTANT, value=0):
    return cv2.copyMakeBorder(image, top, bottom, left, right, border_type, value=value)


def extend_flow(flow, border_size):
    # 扩展后的光流场
    flow_extended = np.zeros((flow.shape[0] + 2 * border_size, flow.shape[1] + 2 * border_size, 2), dtype=np.float32)
    flow_extended[border_size:-border_size, border_size:-border_size, :] = flow

    # 填充上边界
    flow_extended[:border_size, border_size:-border_size, :] = flow[0, :, :]
    # 填充下边界
    flow_extended[-border_size:, border_size:-border_size, :] = flow[-1, :, :]
    # 填充左边界
    flow_extended[border_size:-border_size, :border_size, :] = flow[:, 0, np.newaxis]
    # 填充右边界
    flow_extended[border_size:-border_size, -border_size:, :] = flow[:, -1, np.newaxis]

    # 填充四个角
    flow_extended[:border_size, :border_size, :] = flow[0, 0, :]
    flow_extended[-border_size:, :border_size, :] = flow[-1, 0, :]
    flow_extended[:border_size, -border_size:, :] = flow[0, -1, :]
    flow_extended[-border_size:, -border_size:, :] = flow[-1, -1, :]

    return flow_extended


def apply_optical_flow_with_border(image, flow, border_size=100):
    # 扩展图像边界
    image_extended = extend_border(image, border_size, border_size, border_size, border_size)

    # 扩展光流场
    flow_extended = extend_flow(flow, border_size)

    # 生成变形后的坐标
    h, w = image_extended.shape[:2]
    map_x, map_y = np.meshgrid(np.arange(w), np.arange(h))
    map_x = map_x.astype(np.float32) + flow_extended[:, :, 0]
    map_y = map_y.astype(np.float32) + flow_extended[:, :, 1]

    # 应用光流进行变形
    remapped_image = cv2.remap(image_extended, map_x, map_y, interpolation=cv2.INTER_LINEAR,
                               borderMode=cv2.BORDER_REFLECT101)

    return remapped_image


# 给定图像路径和光流文件路径
image_path_A = r'C:\Users\crxc\Pictures\test\data\4096_image_pair\S11866M3-163_tr2-tc2.png'
image_path_B = r'C:\Users\crxc\Pictures\test\data\4096_image_pair\S11867M1-80_tr2-tc2.png'
flow_path = r'C:\Users\crxc\Pictures\test\data\4096_image_pair\dvf_upsample.npy'
output_path_B = r'C:\Users\crxc\Pictures\test\data\4096_image_pair\output_B.png'
output_path_A = r'C:\Users\crxc\Pictures\test\data\4096_image_pair\output_A.png'

# 加载图像 A 和 B 以及光流场
image_A = cv2.imread(image_path_A)
if image_A is None:
    raise FileNotFoundError(f"无法打开图像文件：{image_path_A}")

image_B = cv2.imread(image_path_B)
if image_B is None:
    raise FileNotFoundError(f"无法打开图像文件：{image_path_B}")

flow = np.load(flow_path)

# 检查光流场尺寸
if flow.shape[0] != image_B.shape[0] or flow.shape[1] != image_B.shape[1]:
    raise ValueError("光流场尺寸与图像 B 尺寸不匹配")

# 设置边界大小
border_size = 256

# 应用光流并保留边界
B_transformed_with_border = apply_optical_flow_with_border(image_B, flow, border_size)

# 保存结果图像 B
cv2.imwrite(output_path_B, B_transformed_with_border)

# 扩展图像 A 的边界
A_extended_with_border = extend_border(image_A, border_size, border_size, border_size, border_size)

# 保存扩展后的图像 A
cv2.imwrite(output_path_A, A_extended_with_border)

print(f"Transformed B image saved to {output_path_B}")
print(f"Extended A image saved to {output_path_A}")
